﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;

namespace EvoBoost.Core
{
	public class SquaredLossFunction : ILossFunction
	{
		public double Calculate(double expectedOutcome, double realOutcome)
		{
			double diff = expectedOutcome - realOutcome;
			return diff * diff;
		}

		private static double[] CalculateResiduals(IList<double> previouslyBoostedOutcomes, IList<double> expectedOutcomes)
		{
			Debug.Assert(previouslyBoostedOutcomes != null);
			Debug.Assert(expectedOutcomes != null);
			Debug.Assert(previouslyBoostedOutcomes.Count == expectedOutcomes.Count);
			Debug.Assert(expectedOutcomes.Count > 0);

			int sampleCount = expectedOutcomes.Count;
			double[] result = new double[sampleCount];
			for (int i = 0; i < sampleCount; ++i)
				result[i] = expectedOutcomes[i] - previouslyBoostedOutcomes[i];

			return result;
		}

		public double Minimize(IList<double> previouslyBoostedOutcomes, IList<double> expectedOutcomes)
		{
			if (previouslyBoostedOutcomes == null) throw new ArgumentNullException("previouslyBoostedOutcomes");
			if (expectedOutcomes == null) throw new ArgumentNullException("expectedOutcomes");
			if (previouslyBoostedOutcomes.Count != expectedOutcomes.Count)
				throw new ArgumentException("Outcome collections should have equal item count.");

			int sampleCount = expectedOutcomes.Count;
			double result = 0;
			for (int i = 0; i < sampleCount; ++i)
				result += expectedOutcomes[i] - previouslyBoostedOutcomes[i];

			return result / sampleCount;
		}

		/// <summary>
		/// Finds optimal split in terms of squared loss for a regression tree node
		/// </summary>
		/// <param name="previouslyBoostedOutcomes">Sum of outcomes from previously boosted trees.</param>
		/// <param name="expectedOutcomes">Expected outcomes for each sample.</param>
		/// <param name="splitterValues">Splitter values for each sample</param>
		/// <remarks>See page 307 of [Hastie, Tibshirani, Friedman 2008] for theoretical background.</remarks>
		/// <returns>Optimal threshold for splitter node.</returns>
		public double FindOptimalSplit(IList<double> previouslyBoostedOutcomes, IList<double> expectedOutcomes, IList<double> splitterValues)
		{
			if (previouslyBoostedOutcomes == null) throw new ArgumentNullException("previouslyBoostedOutcomes");
			if (expectedOutcomes == null) throw new ArgumentNullException("expectedOutcomes");
			if (splitterValues == null) throw new ArgumentNullException("splitterValues");
			if (previouslyBoostedOutcomes.Count != expectedOutcomes.Count || splitterValues.Count != expectedOutcomes.Count)
				throw new ArgumentException("Outcome and splitter value collections should have equal item count.");

			int sampleCount = expectedOutcomes.Count;
			
			// Calculate residual
			double[] residuals = CalculateResiduals(previouslyBoostedOutcomes, expectedOutcomes);

			// Sort splitter values together with residual
			List<SplitterResidualPair> sortedSplitters = new List<SplitterResidualPair>(sampleCount);
			for (int i = 0; i < sampleCount; ++i)
				sortedSplitters.Add(new SplitterResidualPair(splitterValues[i], residuals[i]));
			sortedSplitters.Sort((s1, s2) => Math.Sign(s1.Splitter - s2.Splitter));

			// Select first threshold
			double meanLeft = sortedSplitters.Average(s => s.Residual);
			double meanRight = 0;
			double bestError = sortedSplitters.Sum(s => this.Calculate(s.Residual, meanLeft));
			double bestThreshold = sortedSplitters[sampleCount - 1].Splitter;
			double sumResidualLeft = meanLeft * sampleCount;
			double sumResidualRight = 0;
			double error = bestError;

			// Choose threshold using each splitter value as a candidate
			int index = sampleCount - 1;
			while (index >= 0)
			{
				double meanLeftOld = meanLeft, meanRightOld = meanRight;
				double sumResidualInside = 0;
				int k = 0;
				while(true)
				{
					double residual = sortedSplitters[index].Residual;
					
					// Recalculate mean
					meanRight = (meanRight * (sampleCount - index - 1) + residual) / (sampleCount - index);
					meanLeft = index > 0 ? (meanLeft * (index + 1) - residual) / index : 0;

					// Update residual sums
					sumResidualLeft -= residual;
					sumResidualRight += residual;
					sumResidualInside += residual;
					k += 1;

					if (index == 0 || sortedSplitters[index].Splitter != sortedSplitters[index - 1].Splitter)
						break;
					
					index -= 1;
				}

				index -= 1;

				// Recalculate error using analytic difference formulaes => O(n) instead of O(n^2)
				int n = index + 1;
				double deltaErrorLeft =
				    - 2 * meanLeft * sumResidualLeft + n * meanLeft * meanLeft
				    + 2 * meanLeftOld * (sumResidualLeft + sumResidualInside) - (n + k) * meanLeftOld * meanLeftOld;
				double deltaErrorRight =
				    - 2 * meanRight * sumResidualRight + (sampleCount - n) * meanRight * meanRight
				    + 2 * meanRightOld * (sumResidualRight - sumResidualInside) - (sampleCount - n - k) * meanRightOld * meanRightOld;
				error += deltaErrorLeft + deltaErrorRight;
				
				// Check analytic error update
				//double testError = 0;
				//for (int i = 0; i <= index; ++i)
				//    testError += this.Calculate(sortedSplitters[i].Residual, meanLeft);
				//for (int i = index + 1; i < sampleCount; ++i)
				//    testError += this.Calculate(sortedSplitters[i].Residual, meanRight);
				//Debug.Assert(Math.Abs(testError - error) < 0.00001);
				
				// Compare to the best solution
				if (error < bestError)
				{
					bestError = error;
					bestThreshold = index >= 0 ? sortedSplitters[index].Splitter : sortedSplitters[0].Splitter - 1;
				}
			}

			return bestThreshold;
		}

		public override string ToString()
		{
			return "MSE";
		}

		private struct SplitterResidualPair
		{
			public double Splitter { get; private set; }

			public double Residual { get; private set; }

			public SplitterResidualPair(double splitter, double residual)
				: this()
			{
				this.Splitter = splitter;
				this.Residual = residual;
			}
		}
	}
}
